#include <ros/ros.h>
#include <tf/transform_datatypes.h>

#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/TwistStamped.h>
#include <geometry_msgs/Point32.h>
#include <std_msgs/Float64MultiArray.h>
#include <sensor_msgs/PointCloud2.h>
#include <visualization_msgs/Marker.h>
#include <pcl/common/transforms.h>
#include <pcl/common/angles.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/radius_outlier_removal.h>
#include <iostream>
#include <string>
#include <cmath>
#include <stdlib.h>
#include <mavros_msgs/State.h>
#include <nav_msgs/Odometry.h>
#include <std_msgs/Float64.h>
#include <queue>
#include <cstdint>
#include <trajectory_msgs/JointTrajectory.h>
#include <std_msgs/ColorRGBA.h>
#include <visualization_msgs/MarkerArray.h>
#include <nav_msgs/Path.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>
#include <algorithm>
#include <gmm_voxel_map/gmm.h>
#include "gmap/GMMCommunication.h"

using namespace std;
using namespace Eigen;

int AGENT_ID = 0;
int COLOR_ID = 0;
double init_x = 0;
double init_y = 0;
double init_z = 0;
int count_time = 0; // TEST
double total_time;  // TEST

double motor_yaw = 0; // TODO
double N_SIGMA = 1.6;

ros::Publisher cloud2_pub;
ros::Publisher cluster_centre_pub;
ros::Publisher ellipsoid_pub;
ros::Publisher ellipsoid_marker_resample_pub;
ros::Publisher ellipsoid_marker_part_pub;
ros::Publisher gmm_pub;
ros::Publisher gmm_mix_pub;
ros::Publisher global_position_pub;
// TEST
ros::Publisher pass_cloud_ros_pub;
ros::Publisher Outlier_cloud_ros_pub;
ros::Publisher resample_pc_pub;
ros::Publisher voxel_pc_pub;
ros::Publisher all_pc_pub;
ros::Time time_get_yaw;
ros::Time time_get_pos;

pcl::PointCloud<pcl::PointXYZ>::Ptr cloud0(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud1(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud2(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud3(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_all(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_tmp(new pcl::PointCloud<pcl::PointXYZ>());

double max_height = 5.0;  // TODO
double min_height = 0.54; // TODO
double radius_search = 0.5;
int min_neighbpr = 45;

ros::Timer resample_timer;

gmap::GMMCommunication all_gmm_msg[4];
gvm::GMM gmm;

double PI = 3.1415926;
bool other_gmm_flag = false;

bool gmm_flag[4] = {false, false, false, false};

void PclPassz(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out, double limL, double limH)
{
  pcl::PassThrough<pcl::PointXYZ> pass;
  pass.setInputCloud(cloud_in);
  pass.setFilterFieldName("z"); // 滤波字段设置为z轴方向
  pass.setFilterLimits(limL, limH);
  // 设置保留范围内还是过滤掉范围内, 默认为false,可以注释掉。true为过滤掉范围内的，false为过滤掉范围外的
  //  pass.setFilterLimitsNegative (false);
  pass.filter(*cloud_out);
}

void PclVoxel(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out, double resolution)
{
  pcl::VoxelGrid<pcl::PointXYZ> sorr;
  sorr.setInputCloud(cloud_in);
  sorr.setLeafSize(resolution, resolution, resolution);
  sorr.filter(*cloud_out);
}

void PclOutlier(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out, double RadiusSearch, int MinNeighbors)
{
  pcl::RadiusOutlierRemoval<pcl::PointXYZ> ror; // 创建滤波器对象
  ror.setInputCloud(cloud_in);                  // 设置待滤波点云
  ror.setRadiusSearch(RadiusSearch);            // 设置查询点的半径邻域范围
  ror.setMinNeighborsInRadius(MinNeighbors);    // 设置判断是否为离群点的阈值，即半径内至少包括的点数
  // ror.setNegative(true);						//默认false，保存内点；true，保存滤掉的外点
  ror.filter(*cloud_out);
}

bool PclEuclideanCluster(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, double ClusterTolerance, int MinClusterSize,
                         std::vector<Eigen::Vector3d> &means, std::vector<Eigen::Matrix3d> &var, std::vector<double> &pi_vector, double n_sigma)
{
  pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);
  tree->setInputCloud(cloud_in);
  std::vector<pcl::PointIndices> cluster_indices;
  pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
  ec.setClusterTolerance(ClusterTolerance); // 设置近邻搜索的半径
  ec.setMinClusterSize(MinClusterSize);     // 设置最小聚类点数
  ec.setMaxClusterSize(999999);             // 设置最大聚类点数
  ec.setSearchMethod(tree);
  ec.setInputCloud(cloud_in);
  ec.extract(cluster_indices);

  /// 执行欧式聚类分割，并保存分割结果
  if (cluster_indices.size() == 0) {
    ROS_INFO("\033[1;%dm In PclEuclideanCluster cluster_indices.size()==0 \033[0m", COLOR_ID);
    return false;
  }
  for (std::vector<pcl::PointIndices>::const_iterator it = cluster_indices.begin(); it != cluster_indices.end(); ++it)
  {
    Eigen::Vector3d cluster_centre = Eigen::Vector3d::Zero();
    double num = 0;

    // 求取聚类后的中心
    for (std::vector<int>::const_iterator pit = it->indices.begin(); pit != it->indices.end(); pit++)
    {
      Eigen::Vector3d tmp(cloud_in->points[*pit].x, cloud_in->points[*pit].y, cloud_in->points[*pit].z);
      cluster_centre += tmp;
      num++;
    }
    cluster_centre /= (double)num;
    means.push_back(cluster_centre);
    pi_vector.push_back((double)num);

    Eigen::Matrix3d var_tmp = Eigen::Matrix3d::Zero();

    for (std::vector<int>::const_iterator pit = it->indices.begin(); pit != it->indices.end(); pit++)
    {
      Eigen::Vector3d tmp(cloud_in->points[*pit].x, cloud_in->points[*pit].y, cloud_in->points[*pit].z);
      var_tmp += (tmp - cluster_centre) * (tmp - cluster_centre).transpose();
    }
    var_tmp = n_sigma * var_tmp / (double)num;
    var.push_back(var_tmp);
  }
  return true;
}

/**
 * @param ALL_GMM true ->TOPIC:/gmm_mix_server，混合后的GMM
 * @param ALL_GMM false->TOPIC:/gmm_server，单个GMM
 */
void sendGmmMsg(std::vector<double> pi_vector, std::vector<Eigen::Vector3d> means, std::vector<Eigen::Matrix3d> var, int n_component, bool ALL_GMM)
{
  gmap::GMMCommunication gmm_msg;
  gmm_msg.agent_id = AGENT_ID;
  gmm_msg.gmm_number = n_component;

  for (int i = 0; i < n_component; i++)
  {
    gmm_msg.pi.push_back(pi_vector[i]);
    gmm_msg.means.push_back(means[i](0));
    gmm_msg.means.push_back(means[i](1));
    gmm_msg.means.push_back(means[i](2));

    for (int j = 0; j < 3; j++)
    {
      for (int k = 0; k < 3; k++)
      {
        gmm_msg.vars.push_back(var[i](j, k));
      }
    }
  }
  gmm_msg.header.stamp = ros::Time::now();
  if (ALL_GMM == true) {
    gmm_mix_pub.publish(gmm_msg);
  }
  else {
    gmm_pub.publish(gmm_msg);
    all_gmm_msg[AGENT_ID] = gmm_msg;
  }
}

void visual_ellipsoid(std::vector<Eigen::Vector3d> means, std::vector<Eigen::Matrix3d> var, int id, bool resample_flag)
{
  visualization_msgs::MarkerArray marker_array;
  int cluster_num = means.size();
  // ROS_INFO("means var size %ld %ld",means.size(),var.size());

  for (int i = 0; i < cluster_num; i++)
  {
    for (int j = 0; j < 3; j++)
    {
      for (int k = 0; k < 3; k++)
      {
        if (fabs(var[i](j, k)) < 1e-5)
        {
          var[i](j, k) = 0;
        }
      }
    }

    EigenSolver<Eigen::Matrix3d> es(var[i]);

    Eigen::Matrix3d value_martrix = es.pseudoEigenvalueMatrix();
    Eigen::Matrix3d vector_martrix = es.pseudoEigenvectors();
    vector_martrix.col(2) = vector_martrix.col(0).cross(vector_martrix.col(1));
    // std::cout<<"var: "<<var[i]<<endl;

    // std::cout<<"value_martrix: \n"<<value_martrix<<endl;
    // std::cout<<"vector_martrix: \n"<<vector_martrix<<endl;
    // std::cout<<"means[i]: "<<means[i]<<endl;

    double var0 = pow(value_martrix(0, 0), 0.5);
    double var1 = pow(value_martrix(1, 1), 0.5);
    double var2 = pow(value_martrix(2, 2), 0.5);
    // ROS_INFO("i: %d var: %f %f %f, means: %f %f %f",i,var0,var1,var2,means[i](0),means[i](1),means[i](2));

    visualization_msgs::Marker marker;

    marker.header.frame_id = "map";
    marker.header.stamp = ros::Time();
    marker.id = 23123 + i + id + AGENT_ID;
    marker.type = visualization_msgs::Marker::SPHERE;
    marker.action = visualization_msgs::Marker::ADD;

    marker.lifetime = ros::Duration(1.0);

    marker.pose.position.x = means[i](0);
    marker.pose.position.y = means[i](1);
    marker.pose.position.z = means[i](2);

    Eigen::Quaterniond quaternion(vector_martrix);
    quaternion.normalize();
    marker.pose.orientation.x = quaternion.x();
    marker.pose.orientation.y = quaternion.y();
    marker.pose.orientation.z = quaternion.z();
    marker.pose.orientation.w = quaternion.w();

    marker.scale.x = var0 * 2;
    marker.scale.y = var1 * 2;
    marker.scale.z = var2 * 2;

    marker.color.a = 0.5; // Don't forget to set the alpha!
    marker.color.r = 0.0;
    marker.color.g = 0.0;
    marker.color.b = 0.0;

    marker_array.markers.push_back(marker);
  }
  if (resample_flag) {
    ellipsoid_marker_resample_pub.publish(marker_array);
  }
  else {
    ellipsoid_marker_part_pub.publish(marker_array);
  }

  // visulization of cluster centre
  visualization_msgs::Marker point_marker;

  point_marker.id = 333;
  point_marker.header.frame_id = "map";
  point_marker.header.stamp = ros::Time::now();
  point_marker.type = visualization_msgs::Marker::SPHERE_LIST;
  point_marker.action = visualization_msgs::Marker::ADD;
  point_marker.ns = "lines_and_points";

  point_marker.scale.x = 0.4;
  point_marker.scale.y = 0.4;
  point_marker.scale.z = 0.4;

  point_marker.lifetime = ros::Duration(0);

  point_marker.color.r = 1;
  point_marker.color.g = 0;
  point_marker.color.b = 0;
  point_marker.color.a = 1.0;

  for (unsigned int i = 0; i < means.size(); i++)
  {
    geometry_msgs::Point point;
    point.x = means[i](0);
    point.y = means[i](1);
    point.z = means[i](2);
    point_marker.points.push_back(point);
  }
  cluster_centre_pub.publish(point_marker);
}

void cloudCallback0(const sensor_msgs::PointCloud2ConstPtr &cloud)
{
  if (cloud->data.size() <= 10) return;

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
  pcl::fromROSMsg(*cloud, *cloud_in);

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclPassz(cloud_in, cloud_filter, min_height, max_height); // TODO

  // Remove Outliers
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_second_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclOutlier(cloud_filter, cloud_second_filter, radius_search, min_neighbpr); // TODO

  if (cloud_second_filter->points.size() < 5) {
    ROS_INFO("\033[1;%dm cloud_second_filter < 5 \033[0m", COLOR_ID);
    return;
  }
  cloud0 = cloud_second_filter;
  gmm_flag[0] = true;
}

void cloudCallback1(const sensor_msgs::PointCloud2ConstPtr &cloud)
{
  if (cloud->data.size() <= 10) return;

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
  pcl::fromROSMsg(*cloud, *cloud_in);

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclPassz(cloud_in, cloud_filter, min_height, max_height); // TODO

  // Remove Outliers
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_second_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclOutlier(cloud_filter, cloud_second_filter, radius_search, min_neighbpr); // TODO

  if (cloud_second_filter->points.size() < 5) {
    ROS_INFO("\033[1;%dm cloud_second_filter < 5 \033[0m", COLOR_ID);
    return;
  }
  cloud1 = cloud_second_filter;
  gmm_flag[1] = true;
}

void cloudCallback2(const sensor_msgs::PointCloud2ConstPtr &cloud)
{
  if (cloud->data.size() <= 10) return;

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
  pcl::fromROSMsg(*cloud, *cloud_in);

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclPassz(cloud_in, cloud_filter, min_height, max_height); // TODO

  // Remove Outliers
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_second_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclOutlier(cloud_filter, cloud_second_filter, 0.1, min_neighbpr); // TODO

  if (cloud_second_filter->points.size() < 5) {
    ROS_INFO("\033[1;%dm cloud_second_filter < 5 \033[0m", COLOR_ID);
    return;
  }
  cloud2 = cloud_second_filter;
  gmm_flag[2] = true;
}

void cloudCallback3(const sensor_msgs::PointCloud2ConstPtr &cloud)
{
  if (cloud->data.size() <= 10) return;

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
  pcl::fromROSMsg(*cloud, *cloud_in);

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclPassz(cloud_in, cloud_filter, min_height, max_height); // TODO

  // Remove Outliers
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_second_filter(new pcl::PointCloud<pcl::PointXYZ>());
  PclOutlier(cloud_filter, cloud_second_filter, radius_search, min_neighbpr); // TODO

  if (cloud_second_filter->points.size() < 5) {
    ROS_INFO("\033[1;%dm cloud_second_filter < 5 \033[0m", COLOR_ID);
    return;
  }
  cloud3 = cloud_second_filter;
  gmm_flag[3] = true;
}

void resample_cb(const ros::TimerEvent &e)
{
  // ROS_INFO("\033[1;%dm  resample enter!!!!!!!!!!! \033[0m",COLOR_ID);
  // ros::Time starttime=ros::Time::now();

  bool new_flag = false;
  for (int i = 0; i < 4; i++)
  {
    if (gmm_flag[i] == false) {
      continue;
      ;
    }
    if (i == 0) {
      cloud_tmp = cloud0;
    }
    else if (i == 1) {
      cloud_tmp = cloud1;
    }
    else if (i == 2) {
      cloud_tmp = cloud2;
    }
    else {
      cloud_tmp = cloud3;
    }

    int size_point = cloud_tmp->points.size();
    for (int j = 0; j < size_point; j++) {
      cloud_all->points.push_back(cloud_tmp->points[j]);
    }
    new_flag = true;
    gmm_flag[i] = false;
  }

  if (!new_flag) {
    return;
  }

  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_voxel(new pcl::PointCloud<pcl::PointXYZ>());
  PclVoxel(cloud_all, cloud_voxel, 0.1);
  cloud_all = cloud_voxel;

  sensor_msgs::PointCloud2 all_cloud_ros;
  pcl::toROSMsg(*cloud_all, all_cloud_ros);
  all_cloud_ros.header.stamp = ros::Time::now();
  all_cloud_ros.header.frame_id = "map";
  all_pc_pub.publish(all_cloud_ros);
}

void positionCallback(const geometry_msgs::PoseStamped &msg)
{
  geometry_msgs::PoseStamped global_pos_msg;
  global_pos_msg = msg;
  global_pos_msg.pose.position.x = msg.pose.position.x + init_x;
  global_pos_msg.pose.position.y = msg.pose.position.y + init_y;
  global_pos_msg.pose.position.z = msg.pose.position.z + init_z;

  global_position_pub.publish(global_pos_msg);
}

int main(int argc, char **argv)
{
  ros::init(argc, argv, "gmm_map_node0");
  ros::NodeHandle nh("~");
  srand((unsigned)time(NULL));

  ros::Duration(5.0).sleep();

  nh.getParam("AGENT_ID", AGENT_ID);
  nh.getParam("x", init_x);
  nh.getParam("y", init_y);
  nh.getParam("z", init_z);
  nh.getParam("N_SIGMA", N_SIGMA);

  COLOR_ID = AGENT_ID + 32;

  gmm.setAgentId(AGENT_ID, COLOR_ID, N_SIGMA); // TOASK

  const int LOOPRATE = 30;
  ros::Rate loop_rate(LOOPRATE);

  ros::Subscriber cloud_sub0 = nh.subscribe("/swarm_plan/uav0/convert_pointcloud", 1, cloudCallback0);
  ros::Subscriber cloud_sub1 = nh.subscribe("/swarm_plan/uav1/convert_pointcloud", 1, cloudCallback1);
  ros::Subscriber cloud_sub2 = nh.subscribe("/swarm_plan/uav2/convert_pointcloud", 1, cloudCallback2);
  ros::Subscriber cloud_sub3 = nh.subscribe("/swarm_plan/uav3/convert_pointcloud", 1, cloudCallback3);

  ros::Subscriber position_sub = nh.subscribe("/mavros/local_position/pose", 1, positionCallback);

  cluster_centre_pub = nh.advertise<visualization_msgs::Marker>("/cluster_centr", 1, true);
  ellipsoid_marker_resample_pub = nh.advertise<visualization_msgs::MarkerArray>("/ellipsoid_resample", 1);
  ellipsoid_marker_part_pub = nh.advertise<visualization_msgs::MarkerArray>("/ellipsoid_part", 1);

  gmm_pub = nh.advertise<gmap::GMMCommunication>("/gmm_server", 1, true);         // 发布GMM模型
  gmm_mix_pub = nh.advertise<gmap::GMMCommunication>("/gmm_mix_server", 1, true); // 发布GMM模型
  global_position_pub = nh.advertise<geometry_msgs::PoseStamped>("/global_position", 1, true);

  // TEST
  pass_cloud_ros_pub = nh.advertise<sensor_msgs::PointCloud2>("/pass_pc", 1, true);
  Outlier_cloud_ros_pub = nh.advertise<sensor_msgs::PointCloud2>("/Outlier_pc", 1, true);
  resample_pc_pub = nh.advertise<sensor_msgs::PointCloud2>("/resample_pc", 1, true);
  voxel_pc_pub = nh.advertise<sensor_msgs::PointCloud2>("/voxel_pc", 1, true);

  all_pc_pub = nh.advertise<sensor_msgs::PointCloud2>("/all_pc", 1, true);

  resample_timer = nh.createTimer(ros::Duration(0.05), resample_cb); // TODO

  ros::spin();

  return 0;
}